package simpledb.execution;

import simpledb.common.DbException;
import simpledb.common.Type;
import simpledb.storage.Field;
import simpledb.storage.IntField;
import simpledb.storage.Tuple;
import simpledb.storage.TupleDesc;
import simpledb.transaction.TransactionAbortedException;

import java.util.*;

import static simpledb.execution.Aggregator.Op.AVG;

/**
 * Knows how to compute some aggregate over a set of IntFields.
 */
public class IntegerAggregator implements Aggregator {

    private static final long serialVersionUID = 1L;

    private final int                   gbfield;
    private final Type                  gbfiedltype;
    private final int                   afield;
    private final Op                    what;

    private Map<Field, Integer>         groupMap;
    private Map<Field, Integer>         countMap;
    private Map<Field, List<Integer>>   avgMap;

    /**
     * Aggregate constructor
     * 
     * @param gbfield
     *            the 0-based index of the group-by field in the tuple, or
     *            NO_GROUPING if there is no grouping
     * @param gbfieldtype
     *            the type of the group by field (e.g., Type.INT_TYPE), or null
     *            if there is no grouping
     * @param afield
     *            the 0-based index of the aggregate field in the tuple
     * @param what
     *            the aggregation operator
     */

    public IntegerAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        // some code goes here
        this.gbfield     = gbfield;
        this.gbfiedltype = gbfieldtype;
        this.afield      = afield;
        this.what        = what;

        this.groupMap    = new HashMap<>();
        this.countMap    = new HashMap<>();
        this.avgMap      = new HashMap<>();
    }

    /**
     * Merge a new tuple into the aggregate, grouping as indicated in the
     * constructor
     * 
     * @param tup
     *            the Tuple containing an aggregate field and a group-by field
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        // some code goes here
        IntField afield = (IntField) tup.getField(this.afield);
        Field gbfield = this.gbfield == NO_GROUPING ? null : tup.getField(this.gbfield);

        int newValue = afield.getValue();
        if (gbfield != null && gbfield.getType() != this.gbfiedltype) {
            throw new IllegalArgumentException("Given tuple has wrong type");
        }

        switch (what) {
            case SUM:
                if (!groupMap.containsKey(gbfield))
                    groupMap.put(gbfield, newValue);
                else
                    groupMap.put(gbfield, groupMap.get(gbfield) + newValue);
                break;

            case MIN:
                if (!groupMap.containsKey(gbfield))
                    groupMap.put(gbfield, newValue);
                else
                    groupMap.put(gbfield, Math.min(newValue, groupMap.get(gbfield)));

                break;

            case MAX:
                if (!groupMap.containsKey(gbfield))
                    groupMap.put(gbfield, newValue);
                else
                    groupMap.put(gbfield, Math.max(newValue, groupMap.get(gbfield)));

                break;
            case COUNT:
                if (!groupMap.containsKey(gbfield))
                    groupMap.put(gbfield, 1);
                else
                    groupMap.put(gbfield, groupMap.get(gbfield) + 1);

                break;

            case SC_AVG:
                IntField countField = null;
                if (gbfield == null)
                    countField = (IntField)tup.getField(1);
                else
                    countField = (IntField)tup.getField(2);
                int countValue = countField.getValue();
                if (!this.groupMap.containsKey(gbfield)) {
                    this.groupMap.put(gbfield, newValue);
                    this.countMap.put(gbfield, countValue);
                } else {
                    this.groupMap.put(gbfield, this.groupMap.get(gbfield) + newValue);
                    this.countMap.put(gbfield, this.countMap.get(gbfield) + countValue);
                }

            case SUM_COUNT:
            case AVG:
                if (!avgMap.containsKey(gbfield)) {
                    List<Integer> list = new ArrayList<>();
                    list.add(newValue);
                    avgMap.put(gbfield, list);
                } else {
                    avgMap.get(gbfield).add(newValue);
                }

                break;
            default:
                throw new IllegalArgumentException("Aggregate not supported!");
        }
    }

    /**
     * Create a OpIterator over group aggregate results.
     * 
     * @return a OpIterator whose tuples are the pair (groupVal, aggregateVal)
     *         if using group, or a single (aggregateVal) if no grouping. The
     *         aggregateVal is determined by the type of aggregate specified in
     *         the constructor.
     */
    public OpIterator iterator() {
        // some code goes here
//        throw new
//        UnsupportedOperationException("please implement me for lab2");
        return new IntAggIterator();
    }

    private class IntAggIterator extends AggregateIterator {
        private Iterator<Map.Entry<Field, List<Integer>>> avgIt;
        private boolean isAvg;
        private boolean isSCAvg;
        private boolean isSumCount;

        public IntAggIterator() {
            super(groupMap, gbfiedltype);
            this.isAvg = what.equals(Op.AVG);
            this.isSCAvg = what.equals(Op.SC_AVG);
            this.isSumCount = what.equals(Op.SUM_COUNT);

            if (isSumCount) {
                this.td = new TupleDesc(
                        new Type[] {this.itgbfieldtype, Type.INT_TYPE, Type.INT_TYPE},
                        new String[] {"groupVal", "sumVal", "countVal"});
            }
        }

        @Override
        public void open() throws DbException, TransactionAbortedException {
            super.open();
            if (this.isAvg || this.isSumCount) {
                this.avgIt = avgMap.entrySet().iterator();
            } else {
                this.avgIt = null;
            }
        }

        @Override
        public void close() {
            super.close();
            avgIt = null;
        }

        @Override
        public void rewind() throws DbException, TransactionAbortedException {
            super.rewind();
            if (this.isAvg || this.isSumCount)
                this.avgIt = avgMap.entrySet().iterator();
        }

        @Override
        public boolean hasNext() throws DbException, TransactionAbortedException {
            if (this.isAvg || this.isSumCount) {
                return avgIt.hasNext();
            }
            return super.hasNext();
        }

        @Override
        public Tuple next() throws DbException, TransactionAbortedException, NoSuchElementException {
            Tuple rtn = new Tuple(td);
            if (this.isSumCount || this.isAvg) {
                Map.Entry<Field, List<Integer>> avgOrSumCountEntry = this.avgIt.next();

                Field avgOrSumCountField = avgOrSumCountEntry.getKey();
                List<Integer> avgOrSumCountList = avgOrSumCountEntry.getValue();

                if (this.isAvg) {
                    int value = sumList(avgOrSumCountList) / avgOrSumCountList.size();
                    this.setFields(rtn, value, avgOrSumCountField);
                    return rtn;
                } else {
                    this.setFields(rtn, sumList(avgOrSumCountList), avgOrSumCountField);
                    if (avgOrSumCountField != null) {
                        rtn.setField(2, new IntField(avgOrSumCountList.size()));
                    } else {
                        rtn.setField(1, new IntField(avgOrSumCountList.size()));
                    }

                    return rtn;
                }
            } else if (this.isSCAvg){
                Map.Entry<Field, Integer> entry = this.it.next();
                Field f = entry.getKey();
                this.setFields(rtn, entry.getValue() / countMap.get(f), f);
                return rtn;
            }

            return super.next();
        }

        @Override
        public TupleDesc getTupleDesc() {
            return super.getTupleDesc();
        }

        private int sumList(List<Integer> list) {
            int sum = 0;
            for (int n : list) {
                sum += n;
            }

            return sum;
        }
    }
}
